from flask import Flask, render_template, request, redirect, url_for, session, flash
from flask_login import LoginManager, UserMixin, login_user, logout_user, login_required, current_user
from Flask登录功能验证 import SaltSM3
from database_utils import con_my_sql

# 创建Flask应用实例
app = Flask(__name__)
# 设置用于加密会话数据的密钥
app.secret_key = '123456'

# 初始化Flask-Login
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'index_login'  # 未登录时重定向到登录页面


# 用户类
class User(UserMixin):
    def __init__(self, user_id, username):
        self.id = user_id
        self.username = username


# 用户加载回调函数
@login_manager.user_loader
def load_user(user_id):
    code = "SELECT * FROM users WHERE UserID = %s"
    cursor_ans = con_my_sql(code, (user_id,))

    if cursor_ans:
        user_data = cursor_ans.fetchall()[0]
        return User(user_data['UserID'], user_data['Username'])
    return None


# 定义根路由（'/'）对应的视图函数，用于渲染登录页面（login.html）
@app.route('/')
def index_login():
    return render_template('login.html')


# 定义登出路由（'/logout'）对应的视图函数，从会话中移除用户信息，然后重定向到登录页面
@app.route('/logout')
@login_required
def index_logout():
    logout_user()
    return redirect(url_for('index_login'))


# 定义'/index'路由对应的视图函数，判断会话中是否存在用户，
# 如果存在则渲染主页面（index.html），否则重定向到登录页面
@app.route('/index')
@login_required
def index():
    return render_template('index.html', username=current_user.username)


# 定义'/index_register'路由对应的视图函数，用于渲染注册页面（register.html）
@app.route('/index_register')
def index_register():
    return render_template('register.html')


# 定义处理登录表单提交（POST请求）的路由'/login'对应的视图函数
@app.route('/login', methods=['POST'])
def login():
    name = request.form.get('username')
    pwd = request.form.get('password')

    if not name or not pwd:
        flash('用户名或密码不能为空！')
        return redirect(url_for('index_login'))

    code = "SELECT * FROM users WHERE Username = %s"
    cursor_ans = con_my_sql(code, (name,))
    cursor_select = cursor_ans.fetchall()

    if cursor_select:
        user_data = cursor_select[0]
        password_hash = user_data['PasswordHash']
        salt = user_data['Salt']
        password_salt_hash = SaltSM3.hash_with_salt(pwd, salt)
        if password_salt_hash == password_hash:
            user = User(user_data['UserID'], user_data['Username'])
            login_user(user)
            return redirect(url_for('index'))

    flash('登录失败，请检查您的用户名和密码。')
    return redirect(url_for('index_login'))


# 定义处理注册表单提交（POST请求）的路由'/register'对应的视图函数
@app.route('/register', methods=['POST'])
def register():
    name = request.form.get('username')
    pwd = request.form.get('password')
    confirm_password = request.form.get('confirm_password')

    if not name or not pwd or not confirm_password:
        flash("用户名或密码不能为空")
        return redirect(url_for('index_register'))

    if confirm_password != pwd:
        flash('两次密码输入不一致。')
        return redirect(url_for('index_register'))

    code = "SELECT * FROM users WHERE Username = %s"
    cursor_ans = con_my_sql(code, (name,))
    cursor_select = cursor_ans.fetchone()

    if cursor_select:
        flash('用户名已存在。')
        return redirect(url_for('index_login'))
    else:
        salt = SaltSM3.generate_salt()
        pwd_salt_hash = SaltSM3.hash_with_salt(pwd, salt)
        code = "INSERT INTO `users` (`Username`, `PasswordHash`,`Salt`) VALUES (%s, %s, %s)"
        con_my_sql(code, (name, pwd_salt_hash, salt))
        flash('注册成功！请登录。')
        return redirect(url_for('index_login'))


if __name__ == '__main__':
    # 启动Flask应用，开启调试模式
    app.run(debug=True)


# 测试用户数据
# 张三 123456
# 20221414 20221414
# 1 1